package datastructure.tree;

import java.util.Comparator;

/**
 * @Description: AVL
 * @Author Ammar
 * @Create 2023/4/5 16:34
 */
public class AVLTree<E> extends BalaceBinarySearchTree<E> {

    public AVLTree() {
        this(null);
    }

    public AVLTree(Comparator<E> comparator) {
        super(comparator);
    }

    @Override
    protected TreeNode<E> createNode(E val, TreeNode<E> parent) {
        return new AVLNode<>(val, parent);
    }

    @Override
    protected void afterAdd(TreeNode<E> node) {
        while ((node = node.parent) != null) {
            // 判断 node 节点的祖先节点是否平衡
            if (isBalanced(node)) {
                // 更新高度
                // 在添加节点之后，该节点一定是叶子节点，所以高度为 1
                // 我们需要维护的是祖先节点的高度
                // 节点的高度 = 左右子树的最大高度 + 1
                updateHeight(node);
            } else {
                // 恢复平衡 node 是高度最低的不平衡节点
                rebalance(node);
                // 整棵树恢复平衡
                break;
            }
        }
    }

    @Override
    protected void afterRemove(TreeNode<E> node, TreeNode<E> replacement) {
        while ((node = node.parent) != null) {
            // 判断 node 节点的祖先节点是否平衡
            if (isBalanced(node)) {
                // 更新高度
                updateHeight(node);
            } else {
                // 恢复平衡 node 是高度最低的不平衡节点
                rebalance(node);
            }
        }
    }

    private boolean isBalanced(TreeNode<E> node) {
        return Math.abs(((AVLNode<E>) node).balanceFactor()) <= 1;
    }

    private void updateHeight(TreeNode<E> node) {
        ((AVLNode<E>) node).updateHeight();
    }

    // 统一处理旋转
    private void rebalance(TreeNode<E> grand) {
        TreeNode<E> parent = ((AVLNode<E>) grand).tallerChild();
        TreeNode<E> node = ((AVLNode<E>) parent).tallerChild();
        if (parent.isLeftChild()) { // L
            if (node.isLeftChild()) { // LL
                rotate(grand, node.left, node, node.right, parent, parent.right, grand, grand.right);
            } else { // LR
                rotate(grand, parent.left, parent, node.left, node, node.right, grand, grand.right);
            }
        } else { // R
            if (node.isLeftChild()) { // RL
                rotate(grand, grand.left, grand, node.left, node, node.right, parent, parent.right);
            } else { // RR
                rotate(grand, grand.left, grand, parent.left, parent, node.left, node, node.right);
            }
        }
    }

    @Override
    protected void rotate(
            TreeNode<E> r, // 子树的根节点
            TreeNode<E> a, TreeNode<E> b, TreeNode<E> c,
            TreeNode<E> d,
            TreeNode<E> e, TreeNode<E> f, TreeNode<E> g) {
        super.rotate(r, a, b, c, d, e, f, g);
        // 抽取代码时，需要注意，更新高度的顺序
        // b 和 f 是旋转之前的父节点，需要先更新高度
        // d 是旋转之后的父节点，需要后更新高度
        updateHeight(b);
        updateHeight(f);
        updateHeight(d);
    }


    // 根据 LL RL RR LR 四种情况进行旋转
    // grand: 高度最低的不平衡节点
    private void rebalance1(TreeNode<E> grand) {
        TreeNode<E> parent = ((AVLNode<E>) grand).tallerChild();
        TreeNode<E> node = ((AVLNode<E>) parent).tallerChild();
        if (parent.isLeftChild()) { // L
            if (node.isLeftChild()) { // LL
                rotateRight(grand);
            } else { // LR
                rotateLeft(parent);
                rotateRight(grand);
            }
        } else { // R
            if (node.isLeftChild()) { // RL
                rotateRight(parent);
                rotateLeft(grand);
            } else { // RR
                rotateLeft(grand);
            }
        }
    }

    @Override
    protected void afterRotate(TreeNode<E> grand, TreeNode<E> parent, TreeNode<E> child) {
        super.afterRotate(grand, parent, child);
        // 更新高度
        updateHeight(grand);
        updateHeight(parent);
    }
}
